import numpy as np
from scipy import spatial
import matplotlib.pyplot as plt
from sko.GA import GA_TSP


def cal_total_distance(routine):
    num_points, = routine.shape
    routine = np.concatenate([routine])
    # 计算调整工况的适应度（使用距离计算）
    distance_select_list = [distance_matrix[routine[i], routine[i + 1]] for i in range(num_points - 1)]
    # print(distance_select_list)
    distance_val = sum(distance_select_list)
    return distance_val


if __name__ == '__main__':
    order_num = 5
    order_info = [
        ['A', '001', 5],
        ['A', '002', 5],
        ['B', '003', 7],
        ['B', '004', 7],
        ['C', '005', 6],
    ]
    distance_matrix = np.asarray([
        [0, 1, 3, 3, 2],
        [1, 0, 3, 3, 2],
        [2, 2, 0, 1, 1],
        [2, 2, 1, 0, 1],
        [3, 3, 4, 4, 0],
    ])

    ga_tsp = GA_TSP(func=cal_total_distance, n_dim=order_num, size_pop=50, max_iter=500, prob_mut=1)
    best_points, best_distance = ga_tsp.run()
    print(best_points, best_distance)
